import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import ElasticNet
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

# ==================== 初始化设置 ====================
# 设置中文字体和数学符号显示
plt.rcParams['font.sans-serif'] = ['KaiTi']  # 使用楷体显示中文
plt.rcParams['mathtext.fontset'] = 'stix'    # 数学公式字体
plt.rcParams['axes.unicode_minus'] = False   # 解决负号显示问题

# ==================== 数据准备阶段 ====================
# 设置随机种子保证结果可复现
np.random.seed(42)

# 生成特征数据X：100个样本，1个特征，取值范围[0,2)
X = 2 * np.random.rand(100, 1)

# 生成标签数据：真实模型为 y = 4 + 3X + 噪声
# 4是截距项(bias term)，3是权重系数(weight)
# 噪声服从标准正态分布 N(0,1)
y = 4 + 3 * X + np.random.randn(100, 1)

# 划分训练集和测试集 (80%训练，20%测试)
X_train, X_test, y_train, y_test = train_test_split(
    X, y,
    test_size=0.2,
    random_state=42
)

# ==================== 数据标准化 ====================
# ElasticNet对特征缩放敏感，建议进行标准化
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

# ==================== ElasticNet回归模型 ====================
# 创建ElasticNet回归模型
# alpha: 正则化强度
# l1_ratio: L1正则化比例 (0为纯Ridge，1为纯Lasso)
# max_iter: 最大迭代次数
# tol: 优化容忍度
elastic_net = ElasticNet(
    alpha=0.5,          # 正则化强度
    l1_ratio=0.5,       # L1和L2的混合比例
    max_iter=10000,     # 增加迭代次数确保收敛
    tol=1e-6,           # 更严格的收敛标准
    random_state=42     # 随机种子
)

# 训练模型 (注意y需要ravel()转换为一维数组)
elastic_net.fit(X_train_scaled, y_train.ravel())

# ==================== 模型评估 ====================
# 训练集预测
y_train_pred = elastic_net.predict(X_train_scaled)
train_mse = mean_squared_error(y_train, y_train_pred)

# 测试集预测
y_test_pred = elastic_net.predict(X_test_scaled)
test_mse = mean_squared_error(y_test, y_test_pred)

# ==================== 结果输出 ====================
print("===== ElasticNet回归结果 =====")
print(f"截距项(w0): {elastic_net.intercept_:.4f}")
print(f"系数(w1): {elastic_net.coef_[0]:.4f}")
print(f"训练集MSE: {train_mse:.4f}")
print(f"测试集MSE: {test_mse:.4f}")
print(f"使用的特征数: {np.sum(elastic_net.coef_ != 0)}")

# ==================== 可视化 ====================
plt.figure(figsize=(12, 6))

# 绘制训练数据
plt.scatter(
    X_train, y_train,
    color='blue',
    label='训练数据',
    alpha=0.6
)

# 绘制测试数据
plt.scatter(
    X_test, y_test,
    color='green',
    label='测试数据',
    alpha=0.6
)

# 绘制拟合线
X_plot = np.linspace(0, 2, 100).reshape(-1, 1)
X_plot_scaled = scaler.transform(X_plot)
y_plot = elastic_net.predict(X_plot_scaled)
plt.plot(
    X_plot, y_plot,
    color='red',
    linewidth=2,
    label='ElasticNet拟合线'
)

# 添加图注和标题
plt.title('ElasticNet回归示例 (alpha=0.5, l1_ratio=0.5)', fontsize=14)
plt.xlabel('X特征值', fontsize=12)
plt.ylabel('y标签值', fontsize=12)
plt.legend(fontsize=10)
plt.grid(True)

# 显示图形
plt.show()
